from src.base.agent import Agent
from typing import Dict# 定义专门的agents

# 需求分析agent
requirements_agent = Agent(
    name="Requirements Analysis Agent",
    instructions="You are a requirements analysis specialist. Analyze user requirements and create a detailed specification."
)

# 设计agent
design_agent = Agent(
    name="Design Agent",
    instructions="You are a software design specialist. Create a high-level design based on the requirements specification."
)

# 编码agent
coding_agent = Agent(
    name="Coding Agent",
    instructions="You are a coding specialist. Implement the design in clean, efficient code."
)

# 测试agent
testing_agent = Agent(
    name="Testing Agent",
    instructions="You are a software testing specialist. Review the code, identify potential bugs, and suggest improvements."
)

# 文档agent
documentation_agent = Agent(
    name="Documentation Agent",
    instructions="You are a technical documentation specialist. Create clear and comprehensive documentation for the code."
)

def generate_code(user_requirement: str) -> Dict[str, str]:
    """
    多阶段代码生成函数
    
    参数:
    user_requirement (str): 用户提供的功能需求
    
    返回:
    dict: 包含需求规格、设计文档、代码和文档的字典
    """

    # 阶段1：需求分析
    req_response = requirements_agent.run(
        messages=[{"role": "user", "content": f"Analyze these requirements: {user_requirement}"}]
    )
    requirements_spec = req_response["result"]

    # 阶段2：设计
    design_response = design_agent.run(
        messages=[{"role": "user", "content": f"Create a design based on this specification: {requirements_spec}"}]
    )
    design_doc = design_response["result"]

    # 阶段3：编码
    coding_response = coding_agent.run(
        messages=[{"role": "user", "content": f"Implement this design in code: {design_doc}"}]
    )
    initial_code = coding_response["result"]

    # 阶段4：测试和优化
    testing_response = testing_agent.run(
        messages=[{"role": "user", "content": f"Review and suggest improvements for this code: {initial_code}"}]
    )
    test_feedback = testing_response["result"]

    # 根据测试反馈进行代码优化
    refinement_response = coding_agent.run(
        messages=[
            {"role": "user", "content": f"Refine the code based on this feedback: {test_feedback}"},
            {"role": "user", "content": f"Original code: {initial_code}"}
        ]
    )
    refined_code = refinement_response["result"]

    # 阶段5：文档编写
    doc_response = documentation_agent.run(
        messages=[{"role": "user", "content": f"Create documentation for this code: {refined_code}"}]
    )
    documentation = doc_response["result"]
     
    # 返回包含所有生成内容的字典
    return {
        "requirements": requirements_spec,
        "design": design_doc,
        "code": refined_code,
        "documentation": documentation
    }

# 使用示例
if __name__ == "__main__":
    user_requirement = "Create a Python function that calculates the Fibonacci sequence up to a given number, optimized for performance."
    result = generate_code(user_requirement)

    print("Generated Code:")
    print(result["code"])
    print("\nDocumentation:")
    print(result["documentation"])